"""
@Project    : cosmo-face
@Module     : face_align.py
@Author     : HuangJiWen[huangjiwen@haier.com]
@Created    : 2020/8/24 17:29
@Desc       : 
"""

import os
import argparse

import cv2
import glob
import numpy as np
import torch
from tqdm import tqdm
from skimage import transform as trans

from detection.data import cfg_center_rfb, cfg_center_ghost, cfg_center_resnet18, cfg_center_resnet50
from detection.widerface_test_centerface import FaceDetectionLandmark


def align(img, landmark, **kwargs):
    """人脸对齐"""

    src = np.array([
        [30.2946, 51.6963],
        [65.5318, 51.5014],
        [48.0252, 71.7366],
        [33.5493, 92.3655],
        [62.7299, 92.2041]], dtype=np.float32)

    image_size = kwargs.get('image_size', (112, 112))
    if image_size[1] == 112:
        src[:, 0] += 8.0

    dst = landmark.astype(np.float32)

    trans_form = trans.SimilarityTransform()
    trans_form.estimate(dst, src)
    M = trans_form.params[0:2, :]
    warped = cv2.warpAffine(src=img, M=M, dsize=(image_size[1], image_size[0]), borderValue=0.0)

    return warped


if __name__ == "__main__":
    torch.set_grad_enabled(False)

    parser = argparse.ArgumentParser(description='Centerface')
    parser.add_argument('-m', '--trained_model',
                        default='F:/gitee_project/cosmo-face/detection/weights/resnet50_center_Final.pth',
                        # default='F:/gitee_project/cosmo-face/detection/weights/RFB_center_epoch_245.pth',
                        # default='F:/gitee_project/cosmo-face/detection/weights/resnet_epoch_240.pth',
                        # default='F:/gitee_project/cosmo-face/detection/weights/ghost_0.25_centernet_Final.pth',
                        type=str, help='Trained state_dict file path to open')
    parser.add_argument('--network', default='resnet50', help='Backbone network resnet18 or ghost or RFB or resnet50')
    parser.add_argument('--origin_size', default=True, type=str, help='Whether use origin image size to evaluate')
    parser.add_argument('--cpu', action="store_true", default=False, help='Use cpu inference')
    parser.add_argument('--confidence_threshold', default=0.1, type=float, help='confidence_threshold')
    parser.add_argument('--top_k', default=5000, type=int, help='top_k')
    parser.add_argument('--nms_threshold', default=0.4, type=float, help='nms_threshold')
    parser.add_argument('--keep_top_k', default=750, type=int, help='keep_top_k')
    parser.add_argument('-s', '--save_image', default=True, help='show detection results')
    # parser.add_argument('--vis_thres', default=0.01, type=float, help='visualization_threshold')
    args = parser.parse_args()

    cfg = None
    if args.network == "resnet18":
        cfg = cfg_center_resnet18
    elif args.network == "resnet50":
        cfg = cfg_center_resnet50
    elif args.network == "RFB":
        cfg = cfg_center_rfb
    elif args.network == "ghost":
        cfg = cfg_center_ghost
    else:
        print("Don't support network!")
        exit(0)

    face_detection_landmark = FaceDetectionLandmark(cfg=cfg, args=args)

    # single image predict
    image_path = "../detection/test/temp.jpg"
    img_raw = cv2.imread(filename=image_path, flags=cv2.IMREAD_COLOR)
    img = np.float32(img_raw)
    # 获取关键点
    bounding_boxes, key_points = face_detection_landmark.predict(img)
    # 只取bounding box score最大的人脸做对齐
    max_bounding_box_score_index = bounding_boxes[:, -1].argmax()
    bounding_box = bounding_boxes[max_bounding_box_score_index].reshape(1, 5)
    key_point = key_points[max_bounding_box_score_index].reshape(1, 15)
    # im = face_detection_landmark.draw_results(img, bounding_box, key_point)
    # cv2.imwrite("../detection/output/center_resnet50_test.jpg", im)
    five_points = key_point[0, :10].reshape(-1, 2)
    # 人脸对齐
    align_img = align(img=img_raw, landmark=five_points, image_size=(112, 112))
    cv2.imwrite("../detection/output/align_image_112.jpg", align_img)

    print("***************************************************")
    folder_test = True
    if folder_test:
        image_path_lst = glob.glob(pathname="../recognition/data/faces_database/data/*/*.jpg")
        # image_path_lst = glob.glob(pathname="./data/*.jpg")
        for image_path in tqdm(image_path_lst):
            img_raw = cv2.imread(filename=image_path, flags=cv2.IMREAD_COLOR)
            img = np.float32(img_raw)
            # 获取关键点
            bounding_boxes, key_points = face_detection_landmark.predict(img)
            # 只取bounding box score最大的人脸做对齐
            max_bounding_box_score_index = bounding_boxes[:, -1].argmax()
            bounding_box = bounding_boxes[max_bounding_box_score_index].reshape(1, 5)
            key_point = key_points[max_bounding_box_score_index].reshape(1, 15)
            five_points = key_point[0, :10].reshape(-1, 2)
            # 人脸对齐
            align_img = align(img=img_raw, landmark=five_points)

            save_folder = os.path.dirname(image_path).replace("faces_database/data", "faces_database/align_data")
            if not os.path.exists(save_folder):
                os.makedirs(save_folder)
            file_name = os.path.basename(image_path)
            cv2.imwrite(save_folder + "/" + file_name, align_img)
